import pysam
import argparse
import pod5
import numpy as np

def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument('--inbam', type=str, default="/public/data/wangjianxin/nipeng/workspace/arab_col0_csu2_nextomics.ont_r10.4.1_simplex.pod5_pass.bam", required=False,
                        help="bam ready to be split")
    parser.add_argument('--outbam', type=str, default="/public/home/hpc234701005/data/R10.4/arab/train/arab_col0_train.bam", required=False,
                        help="where the splited bam to store")
    return parser.parse_args()


def get_refloc_of_methysite_in_motif(seqstr, motifset, methyloc_in_motif=0):
    """

    :param seqstr:
    :param motifset:
    :param methyloc_in_motif: 0-based
    :return:
    """
    motifset = set(motifset)
    strlen = len(seqstr)
    motiflen = len(list(motifset)[0])
    sites = []
    for i in range(0, strlen - motiflen + 1):
        if seqstr[i : i + motiflen] in motifset:
            sites.append(i + methyloc_in_motif)
    return sites

def deal_bam(inbam,outbam,motif_seqs,mod_loc):
    input_bam = pysam.AlignmentFile(inbam, 'rb',check_sq=False)
    output_bam = pysam.AlignmentFile(outbam, 'wb', header=input_bam.header)

    for read in input_bam:
        
        read_dict=read.to_dict()
        #mv_table=read_dict["mv"]
        # strand_code = 0 if read.is_reverse else 1
                
        # strand = "-" if strand_code==0 else "+"
        # ref_seq=read.get_reference_sequence()
        # if strand=="-":
        #     ref_seq = complement_seq(ref_seq)   
        # r_to_q_poss = parse_cigar(cigartuple, strand_code, len(ref_seq))
        seq=read.get_forward_sequence()
        offsets = get_refloc_of_methysite_in_motif(seq, set(motif_seqs), mod_loc)
        oread=pysam.AlignedSegment.from_dict(read_dict,header=input_bam.header)
        output_bam.write(oread)
    
    input_bam.close()
    output_bam.close()


if __name__ == '__main__':
    args=parse_args()
    deal_bam(args.inbam,args.outbam)
